#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   Speed.py
@Time    :   2024-05-29 09:50
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   当前文件作用
"""
import logging
import sys
from argparse import ArgumentParser
from concurrent.futures import ThreadPoolExecutor

import requests

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)

# 文件URL
url = 'http://speedtest.zju.edu.cn/1000M'
# 请求头
headers = {
	'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
	'Accept-Language': 'zh-CN,zh;q=0.9',
	'Connection': 'keep-alive',
	# 'Referer': 'http://speedtest.zju.edu.cn/',
	'Upgrade-Insecure-Requests': '1',
	'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.6045.160 Safari/537.36'
}


class Speed:
	def __init__(self, urls, threads=0):
		"""

		:param urls:
		"""
		self.threads = threads
		self.url = urls  # 文件URL
		self.num_threads = 0  # 线程数
		self.chunk_size = 0  # 每块大小
		# 获取文件大小
		self.total_size = 0

	# 计算每个线程下载的块大小
	def calculate_chunk_size(self):
		"""
		计算每个线程下载的块大小
		:return:
		"""
		self.chunk_size = self.total_size // self.num_threads
		self.chunk_size = 2048 * 1024 * 1024

	def thread_download(self):
		"""

		:return:
		"""
		# 打印下载参数
		logger.info(f"Downloading {self.url} with {self.num_threads} threads")
		# 打印文件大小
		logger.info(f"File size: {self.total_size} bytes")
		# 打印块大小
		logger.info(f"Chunk size: {self.chunk_size} bytes")

		# 使用线程池下载文件
		with ThreadPoolExecutor(max_workers=self.num_threads) as executor:
			# 计算每个线程的工作范围
			futures = []
			for i in range(self.num_threads):
				start = i * self.chunk_size
				# 确保不会超出文件大小
				end = min((i + 1) * self.chunk_size - 1, self.total_size - 1)
				futures.append(executor.submit(self.download_chunk, start, end))

			# 等待所有线程完成
			for future in futures:
				future.result()

	def download_chunk(self, start, end):
		"""

		:param start: 开始位置
		:param end: 结束位置
		:return:
		"""
		range_header = {'Range': f'bytes={start}-{end}'}
		headers.update(range_header)
		try:
			response = requests.get(url, headers=headers, stream=True)
			response.raise_for_status()  # 确保请求成功
			downloaded = 0
			for chunk in response.iter_content(chunk_size=self.chunk_size):  # 1KB的块大小用于迭代
				downloaded += len(chunk)
				# 立即释放内存
				del chunk

		except requests.exceptions.RequestException as e:
			logger.error(f"Error downloading chunk: {e}")
			return
		logger.info(f"Downloaded {downloaded} bytes from chunk {start}-{end}")

	def get_size(self):
		"""
		获取文件大小
		:return:
		"""
		response = requests.head(self.url, headers=headers)
		self.total_size = int(response.headers['Content-Length'])

	# 通过线程数量及文件大小分配合适的线程数
	def get_num_threads(self):
		if 500 * 1024 * 1024 < self.total_size < 1024 * 1024 * 1024:  # 如果文件大小小于1GB
			return 4
		elif 1024 * 1024 * 1024 < self.total_size < 1024 * 1024 * 1024 * 10:  # 如果文件大小小于10GB但大于1GB
			return 10
		else:
			return 10

	def is_threads_available(self):
		"""
		判断是否支持多线程下载
		:return:
		"""
		try:
			response = requests.head(self.url, headers=headers)
			if 'Accept-Ranges' in response.headers and response.headers['Accept-Ranges'] == 'bytes':
				logger.info("This file supports multi-threaded downloads")
				return True
			else:
				logger.warning("This file does not support multi-threaded downloads")
				return False
		except requests.exceptions.RequestException:
			logger.error("An error occurred while checking the file")
			return False

	def start(self):
		self.get_size()  # 获取文件大小
		if self.is_threads_available():
			if self.threads != 0:
				logger.info(f"Using {self.threads} threads")
				self.num_threads = self.threads
			else:
				logger.info("Automatically determining the number of threads")
				self.num_threads = self.get_num_threads()  # 获取线程数
		else:
			logger.warning("This file does not support multi-threaded downloads")
			self.num_threads = 1
		self.calculate_chunk_size()  # 计算每个线程下载的块大小
		self.thread_download()


if __name__ == '__main__':
	# 总文件大小（以字节为单位），假设为1GB，需要根据实际情况调整
	total_size = 1000 * 1024 * 1024
	# 每个线程下载的块大小（以字节为单位）
	chunk_size = 10 * 1024 * 1024  # 10MB
	arg = ArgumentParser(description='当前脚本版本: 1.0', prog="网速持续测速脚本")
	h = f"设置URL，例如：http://speedtest.zju.edu.cn/1000M"
	arg.add_argument('-u', '--url', type=str, help=h, required=False)
	# 传入线程数量,默认10:
	h = f"设置线程数量，默认自动根据文件大小自动分配"
	arg.add_argument('-t', '--threads', type=int, help=h, required=False, default=0)
	args = arg.parse_args()
	arg_url = args.url
	# 获取线程
	arg_threads = args.threads
	# 如果没有输入URL,则终止
	if arg_url is None:
		print("请传入URL")
		exit(1)
	td = Speed(urls=arg_url, threads=arg_threads)
	td.start()
